#!/usr/bin/env python
'''
Given a fasta file of Aceview isoforms, the program
outputs one gene one file which contains all the its isoforms
'''
from get_unique_AceviewGene_mappings import *
from common import output_all_lines
from common import isFileExist
import os.path

def _filename_lize_( filename ):
    filename = filename.replace( '-', '_' )
    filename = filename.replace( ':', '_' )
    return( filename )
#    returnStr = ""
#    valid_chars = "_.%s%s" % ( string.ascii_letters, string.digits )
#    for c in filename:
#        if ( c in valid_chars ):
#            returnStr += c
#        else:
#            returnStr += '_'
#    return returnStr

def get_isoforms4each_AceView_gene( fasta_file ):
    ''' Given fasta file, output one gene per file,
     each file contains all its isofroms name and length '''

    ifile = open( fasta_file, mode = "r" )
    # read until the first sequence header
    gene_name = ""
    isoform_name = ""
    isoform_length = 0
    isoforms_lines = [] # set to collect junction info per contig  
    def add_isoform_output_line():
        if( isoform_length > 0 ):
            output_line = isoform_name + '\t' + str( isoform_length ) + '\n'
            output_line = _filename_lize_( output_line )
            isoforms_lines.append( output_line )

    while 1:
        line = ifile.readline()
        # If file is end, add the last isoform
        # output isoforms of the last gene 
        if ( not line ) or ( line == '\n' ):
            add_isoform_output_line()
            output_all_lines( gene_name + ".isoforms", isoforms_lines )
            break;
        # read fasta and get isoform name and length
        if ( line[0] == '>' ):
            # end of last isoform; add the isoform  
            add_isoform_output_line()
            # new isoform 
            isoform_length = 0
            isoform_name = _filename_lize_( line[1::].strip() )
            current_gene_name = extract_AceView_gene_name( isoform_name )
            # check it belongs to previous gene 
            if ( current_gene_name != gene_name ):
                output_all_lines( gene_name + ".isoforms", isoforms_lines )
                del isoforms_lines[:]
                gene_name = current_gene_name
        else :
            isoform_length += ( len( line ) - 1 )
    ifile.close()

def print_isoforms_and_their_lengths( fasta_file, gene_name, outFN ):
    ''' Given a fasta file and a gene, output a file containing the
    isoforms and their length corresponding to the gene. '''
    print "Collect isoforms of " + gene_name
    isFileExist( fasta_file )
    ifile = open( fasta_file, mode = "r" )
    # read until the first sequence header
    current_gene_name = ""
    isoform_name = ""
    isoform_length = 0
    isoforms_lines = [] # set to collect junction info per contig  
    def add_isoform_output_line_with_gene_name( current_gene_name ):
        if( isoform_length > 0 and gene_name in current_gene_name ):
            output_line = isoform_name + '\t' + str( isoform_length ) + '\n'
            output_line = _filename_lize_( output_line )
            isoforms_lines.append( output_line )

    #sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    while 1:
        line = ifile.readline()
        # If file is end, add the last isoform        
        # output isoforms of the last gene 
        if ( not line ):
            add_isoform_output_line_with_gene_name( current_gene_name )
            output_all_lines( outFN, isoforms_lines )
            print "There are total " + str( len( isoforms_lines ) )\
             + " isoforms\n";
            break;
        elif ( line == '\n' ):
            continue

        # read fasta and get isoform name and length
        if ( line[0] == '>' ):
            # end of last isoform; add the isoform  
            add_isoform_output_line_with_gene_name( current_gene_name )
            # new isoform 
            isoform_length = 0
            header = line[1::].strip() # be careful about the isoform name 
            isoform_name = _filename_lize_( header.split()[0] )
            current_gene_name = extract_AceView_gene_name( isoform_name )
            # print current_gene_name + "\n"
        else :
            isoform_length += ( len( line ) - 1 )
    ifile.close()

def get_Isoforms( gene_name ):
    ''' return a dictionary for isoform names and lengths
    '''
    isofroms = []
    ifile = open( gene_name + ".isoforms", mode = "r" )
    for line in ifile:
        info = line.split()
        isofroms.append( ( info[0], int( info[1] ) ) )
    return( dict( isofroms ) )

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print __doc__
    print "The first argv is the fasta file with AceView's isoform"
    print "The second argv is the sub-strings used to extract isoformss"
    print "The third argv is the output file"

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error, msg:
        print msg
        sys.exit( 2 );
    #process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            print __doc__
            sys.exit();

    if( len( args ) == 1 or len( args ) == 3 ):
        ext_name = os.path.splitext( args[0] )
        if ( ext_name[1] == '.fa' or ext_name[1] == '.fasta' ):
            if len( args ) == 1:
                get_isoforms4each_AceView_gene( args[0] )
            else:
                fastFN = args[0]
                geneSubStr = args[1]
                outFN = args[2]
                print_isoforms_and_their_lengths( fastFN, geneSubStr, outFN )
        else:
            usage()
if __name__ == '__main__':
    main()

